/*
 * COPYRIGHT:	2010 Christopher Smith, licenced under the GNU GPL v2; see COPYING in the top level directory
 * PROJECT:		sdedit
 * FILE:		privileges.cpp
 * PURPOSE:		token privilege control and helper functions
 *				
 */

//the access rights one is entitled to by holding various privileges
ACCESS_MASK StandardPrivilegedRights(UCHAR Privileges)
{
	ACCESS_MASK Ret = 0;

	if(Privileges & PRIV_TAKE_OWN)
		Ret |= WRITE_OWNER;

	if(Privileges & PRIV_SECURITY)
		Ret |= ACCESS_SYSTEM_SECURITY;

	return Ret;
}

//the security related access rights one is entitled to by holding various privileges on a storage resource (e.g. file or key)
ACCESS_MASK StoragePrivilegedSecurityRights(UCHAR Privileges)
{
	ACCESS_MASK Ret = StandardPrivilegedRights(Privileges);

	if(Privileges & PRIV_BACKUP)
		Ret |= READ_CONTROL | ACCESS_SYSTEM_SECURITY;

	if(Privileges & PRIV_RESTORE)
		Ret |= WRITE_DAC | WRITE_OWNER | ACCESS_SYSTEM_SECURITY;
	
	return Ret;
}

//all the access rights one is entitled to by holding various privileges on a storage resource (e.g. file or key)
ACCESS_MASK StoragePrivilegedRights(UCHAR Privileges)
{
	ACCESS_MASK Ret = StoragePrivilegedSecurityRights(Privileges);

	if(Privileges & PRIV_BACKUP)
		Ret |= GENERIC_READ;

	if(Privileges & PRIV_RESTORE)
		Ret |= GENERIC_WRITE;
	
	return Ret;
}

SDE_ENUM_VAL PrivUsePrivEnumVals[] = 
{
	{L"Backup", PRIV_BACKUP},
	{L"Restore", PRIV_RESTORE},
	{L"TakeOwn", PRIV_TAKE_OWN},
	{L"Security", PRIV_SECURITY},
	{L"Relabel", PRIV_RELABEL}
};

SDE_ENUM PrivUsePrivEnum = 
{
	RTL_NUMBER_OF(PrivUsePrivEnumVals),
	PRIV_BACKUP | PRIV_RESTORE | PRIV_TAKE_OWN | PRIV_SECURITY | PRIV_RELABEL,
	PrivUsePrivEnumVals
};

SDE_OPT PrivOptions[] =
{
	{L"Use", SDE_OPT_SET, SptMultiEnumOr, &PrivUsePrivEnum},
};

enum PrivUseOptionIndex //synch with PrivUseOptions
{
	PuoiUse,
};

SDE_OPTS PrivOpts =
{
	RTL_NUMBER_OF(PrivOptions),
	PrivOptions
};

template <ULONG Count> struct TOKEN_PRIVILEGES_AND_ARRAY
{
	union
	{
		struct
		{
			ULONG PrivilegeCount;
			LUID_AND_ATTRIBUTES Privileges[Count];
		};
		TOKEN_PRIVILEGES TokenPrivileges;
	};

	TOKEN_PRIVILEGES_AND_ARRAY() :
		PrivilegeCount(Count)
	{}
};

template <ULONG Count> struct PRIVILEGE_SET_AND_ARRAY
{
	union
	{
		struct
		{
			ULONG PrivilegeCount;
			ULONG Control;
			LUID_AND_ATTRIBUTES Privileges[Count];
		};
		PRIVILEGE_SET PrivilegeSet;
	};

	PRIVILEGE_SET_AND_ARRAY() :
		PrivilegeCount(Count)
	{}
};

//ModifyPrivs attempt to enable privileges according to specified options when created
//and returns those privileges to their previous state when destroyed
class ModifyPrivs
{
	HANDLE Token;
	union
	{
		TOKEN_PRIVILEGES PreviousState;
		UCHAR PreviousStateBytes[1024];
	};

public:
	UCHAR Privileges;

	ModifyPrivs(SDE_OPTS *Opts) :
		Token(NULL),
		Privileges(0)
	{
		OpText Ot(L"Adjusting user privileges");

		SDE_OPT *Opt = Opts->Opts;

		OpenThreadToken(GetCurrentThread(), TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, FALSE, &Token);
			
		if(!Token)//if not thread, process
			RaiseOnWin32ErrorT(OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, &Token));

		ULONG ReqLen;
		TOKEN_PRIVILEGES_AND_ARRAY<5> PrivStart;

		RaiseOnWin32ErrorT(LookupPrivilegeValue(NULL, SE_BACKUP_NAME, &PrivStart.Privileges[0].Luid));
		RaiseOnWin32ErrorT(LookupPrivilegeValue(NULL, SE_RESTORE_NAME, &PrivStart.Privileges[1].Luid));
		RaiseOnWin32ErrorT(LookupPrivilegeValue(NULL, SE_TAKE_OWNERSHIP_NAME, &PrivStart.Privileges[2].Luid));
		RaiseOnWin32ErrorT(LookupPrivilegeValue(NULL, SE_SECURITY_NAME, &PrivStart.Privileges[3].Luid));
		
		//Relabel doesn't exist prior to NT6. Silently ignore if OS doesn't support it
		if(!LookupPrivilegeValue(NULL, L"SeRelabelPrivilege", &PrivStart.Privileges[4].Luid))
			PrivStart.PrivilegeCount--;

		PrivStart.Privileges[0].Attributes = Opt[PuoiUse].Enum->Choice & PRIV_BACKUP ? SE_PRIVILEGE_ENABLED : 0;
		PrivStart.Privileges[1].Attributes = Opt[PuoiUse].Enum->Choice & PRIV_RESTORE ? SE_PRIVILEGE_ENABLED : 0;
		PrivStart.Privileges[2].Attributes = Opt[PuoiUse].Enum->Choice & PRIV_TAKE_OWN  ? SE_PRIVILEGE_ENABLED : 0;
		PrivStart.Privileges[3].Attributes = Opt[PuoiUse].Enum->Choice & PRIV_SECURITY ? SE_PRIVILEGE_ENABLED : 0;
		PrivStart.Privileges[4].Attributes = Opt[PuoiUse].Enum->Choice & PRIV_RELABEL ? SE_PRIVILEGE_ENABLED : 0;

		BOOL Result = AdjustTokenPrivileges(Token, FALSE, &PrivStart.TokenPrivileges, 1024, &PreviousState, &ReqLen);
		
		if(!Result)
		{
			ULONG Status = GetLastError();

			if(Status != ERROR_NOT_ALL_ASSIGNED)
				RaiseWin32Error(Status);
		}

		PRIVILEGE_SET_AND_ARRAY<5> Pset;
		Pset.Control = 0;
		Pset.Privileges[0].Attributes = 0;
		Pset.Privileges[1].Attributes = 0;
		Pset.Privileges[2].Attributes = 0;
		Pset.Privileges[3].Attributes = 0;
		Pset.Privileges[4].Attributes = 0;

		RtlCopyMemory(Pset.Privileges, PrivStart.Privileges, sizeof(LUID_AND_ATTRIBUTES) * 5); //no need to lookup privs again

		RaiseOnWin32ErrorT(PrivilegeCheck(Token, &Pset.PrivilegeSet, &Result));

		if(Pset.Privileges[0].Attributes & SE_PRIVILEGE_USED_FOR_ACCESS)
			Privileges |= PRIV_BACKUP;

		if(Pset.Privileges[1].Attributes & SE_PRIVILEGE_USED_FOR_ACCESS)
			Privileges |= PRIV_RESTORE;

		if(Pset.Privileges[2].Attributes & SE_PRIVILEGE_USED_FOR_ACCESS)
			Privileges |= PRIV_TAKE_OWN;

		if(Pset.Privileges[3].Attributes & SE_PRIVILEGE_USED_FOR_ACCESS)
			Privileges |= PRIV_SECURITY;

		if(Pset.Privileges[4].Attributes & SE_PRIVILEGE_USED_FOR_ACCESS)
			Privileges |= PRIV_RELABEL;
		
		Ot;
	}

	~ModifyPrivs()
	{
		if(Token)
		{
			AdjustTokenPrivileges(Token, FALSE, &PreviousState, 0, NULL, NULL);
			CloseHandle(Token);
		}
	}
};